
import glob
import os
import math

import cv2
import numpy as np
import PIL
import json

from labelme import utils as labelme_utils
import matplotlib.pyplot as plt

# 1----------------------------------------------
def polygon2box(polygon): # get_box_frm_mask
    """轮廓转为 框（中心点xy，宽，长）
    
    polygons=contour=points: 点的数组
    """
    rows = polygon[:, 0]
    clos = polygon[:, 1]

    # 左上角
    left_top_r = np.min(rows)  # x
    left_top_c = np.min(clos)  # y

    # 右下角
    right_bottom_r = np.max(rows)
    right_bottom_c = np.max(clos)

    # 中心点
    x = int(0.5 * (left_top_r + right_bottom_r))
    y = int(0.5 * (left_top_c + right_bottom_c))

    # 宽，长
    w = right_bottom_r-left_top_r
    h = right_bottom_c-left_top_c

    return x, y,w, h

def polygon2rotatedbox(points): # get_rotatedbox_frm_mask
    return cv2.minAreaRect(points)

def polygon2mask(img_shape, polygon): # polygons_to_mask
    mask = np.zeros(img_shape, dtype=np.uint8)
    mask = PIL.Image.fromarray(mask)
    xy = list(map(tuple, polygon))
    PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=np.uint8)
    return mask



# 2----------------------------------------------

def parse_json_file(path_json):
    img = None
    masks = []
    rbboxes = []
    labels = []
        
    with open(path_json,'r', encoding='UTF-8') as fp:
        json_data = json.load(fp) 
        img = labelme_utils.img_b64_to_arr(json_data['imageData'])
        height, width = img.shape[:2]
            
        for shapes in json_data['shapes']:
            points=np.array(shapes['points'],np.float32)
            label = int(shapes['label']) if('shapes' in json_data) else 1
            labels.append(label)
                
            mask = polygon2mask([height,width], points)
            mask = np.uint8(mask)#np.uint8(morphology.binary_dilation(mask, morphology.disk(1)))
            masks.append(mask) 
            
            # rbbox = polygon2box(points) 
            rbbox = polygon2rotatedbox(points)
            rbboxes.append(np.array(rbbox))
        
    return img, masks, rbboxes,labels

# 3----------------------------------------------

def img_mask2masked(img,mask):
    masked = cv2.add(img, np.zeros(np.shape(img), dtype=np.uint8), mask=mask)
    masked[masked==0]=255
    return masked

def get_corner(rbbox,shape):
    H,W = shape
    
    x, y,w,h = rbbox
    # 左上角
    x1 = int(max(0, x-0.5*w))
    y1 = int(max(0, y-0.5*h))
    # 右下角
    x2 = int(min(W, x+0.5*w))
    y2 = int(min(H, y+0.5*h))
    
    return (x1,y1),(x2,y2)



def crop1(img,rbbox): # 不旋转
    (x1,y1),(x2,y2) = get_corner(rbbox,img.shape[:2])

    print('crop1_1',rbbox,img.shape)
    print('crop1_2',(x1,y1),(x2,y2))
    return img[y1:y2,x1:x2]

def crop2(img,rbbox): # 旋转
    (center, size,angle) = rbbox
    print('crop2',rbbox)

    center = tuple(map(int,center))
    size   = (int(size[1]),int(size[0]))

    M = cv2.getRotationMatrix2D(center, angle+90, 1)
    img_rot = cv2.warpAffine(img, M, (img.shape[1],img.shape[0]))
    img_crop = cv2.getRectSubPix(img_rot, size, center)
    print('crop2',center,size,angle)

    return img_crop

def crop(img,rbbox):
    return crop1(img,rbbox) if len(rbbox)==4 else crop2(img,rbbox)


def save_img(img,path,label,path_json,index):
    if label==23:
        label = 'X'
    if label==24:
        label = 'Y'

    path_dir2 = f'{path}/{label}'
    if not os.path.exists(path_dir2):
        os.mkdir(path_dir2)
    
    img_name = os.path.basename(path_json)[:-5] # 去掉 .json
    path_img = f'{path_dir2}/{img_name}__{index}.png'
    cv2.imwrite(path_img,img)
    print('filename',path_img)

def rotate_ornot(img):
    shape = img.shape
    if shape[0] < shape[1]:
        img = cv2.rotate(img,cv2.ROTATE_90_CLOCKWISE)
    return img

# 4----------------------------------------------
def main(path_dir,target_dir):
    path_json_s = glob.glob(f'{path_dir}/*/*.json') # 获得所有json文件
    print(path_json_s)
    for i,path_json in enumerate(path_json_s):
        (img, masks, rbboxes,labels) = parse_json_file(path_json)
        for index,(mask,rbbox,label) in enumerate(zip(masks,rbboxes,labels)):
            if label==-1:continue
            masked = img_mask2masked(img,mask)
            # plt.imshow(masked)
            croped = crop(masked,rbbox)
            rotated = rotate_ornot(croped)
            gray = cv2.cvtColor(rotated,cv2.COLOR_RGB2GRAY)
            RGB = cv2.cvtColor(gray,cv2.COLOR_GRAY2BGR)
            save_img(RGB,target_dir,label,path_json,index)
            
            # if i==0:
            #     plt.imshow(croped)
            #     return
        
# 5----------------------------------------------

path_dir = 'mydata/sample'
target_dir = 'mydata/test'
main(path_dir,target_dir)

# 6----------------------------------------------



# 7----------------------------------------------



# 8----------------------------------------------



# 9----------------------------------------------